/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#ifndef LINAL_DEVICE_GPU_HXX
#define LINAL_DEVICE_GPU_HXX

namespace linal {
  namespace gpu {

    typedef struct FLA_Obj_view FLA_Obj;
    typedef class  Matrix_*     Matrix;
    typedef class  Flat_*       Flat;
    typedef class  Device_Obj_* Device_Obj;
    typedef class  Flat_Obj_*   Flat_Obj;
    typedef class  Manager_*    Manager;
    
    // -------------------------------------------------------------- 
    // ** GPU device manager
    class Manager_ : public Device_,
                     public Cache_<Flat_Obj_> {
    private:
    protected:
    public:
      Manager_() : Device_(0, LINAL_DEVICE_GPU, LINAL_DEVICE_GPU), Cache_<Flat_Obj_>() { }  
      Manager_(Int_ id, Int_ group) : Device_(id, group, LINAL_DEVICE_GPU), Cache_<Flat_Obj_>() { }  
      virtual ~Manager_() { }

      /*
      virtual std::ostream& content(std::ostream& os) const {
        Device_::content(os);
        Cache_<Flat_Obj_>::content(os);
        return os;
      }
      */

      inline void set_device() {
        cudaError_t error = cudaSetDevice( this->_id );
        LINAL_ERROR(error == cudaSuccess, ">> cudaSetDevice Error");
      }
      inline void valid_device() {
        Int_ device=0;
        cudaError_t error = cudaGetDevice( &device );
        LINAL_ERROR(error == cudaSuccess, ">> cudaGetDevice Error");
        LINAL_ERROR(device == this->_id, ">> Device is not associated correctly");
      }

      inline virtual void initialize( Int_ balancer, 
                                      Int_ cache, Int_ n_bin, Int_ n_work ) {

        this->Device_::initialize( balancer, cache );
	this->set_device();
	this->valid_device();
        this->Cache_<Flat_Obj_>::initialize( n_bin, n_work );
      }
      
      inline virtual void finalize() {
	this->set_device();
        this->Cache_<Flat_Obj_>::finalize();

        this->Device_::finalize();
      }

      inline virtual void critical_begin() { 
	omp_set_lock  ( &(this->_omp) ); 
      }
      inline virtual void critical_end()   { 
	omp_unset_lock( &(this->_omp) ); 
      }
      inline virtual void invalidate(FLA_Obj &A) {
        this->pre_sync_obj_in_bin(A);
      }

      
      // devices are pre-selected by the mapping policy 2d heuristic distribution 
      // or 1d distribution.
      virtual Int_ chol_nopiv  ( Int_ uplo, FLA_Obj &A ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }
      virtual Int_ chol_piv    ( Int_ uplo, FLA_Obj &A, FLA_Obj &p ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }
      virtual Int_ ldl_nopiv   ( Int_ uplo, FLA_Obj &A ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }
      virtual Int_ ldl_piv     ( Int_ uplo, FLA_Obj &A, FLA_Obj &p ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }

      virtual Int_ lu_nopiv    ( FLA_Obj &A ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }
      virtual Int_ lu_piv      ( FLA_Obj &A, FLA_Obj &p ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }
    
      virtual Int_ qr          ( FLA_Obj &A, FLA_Obj &T ) {
        return LINAL_DEVICE_RETURN_NOT_SUPPORT;
      }

      virtual Int_ gemm( Int_ transa,  Int_ transb,
                        FLA_Obj &alpha, FLA_Obj &A,  FLA_Obj &B,
                        FLA_Obj &beta,  FLA_Obj &C ) {

        this->gemm_stat();
	this->set_device();

        Flat_Obj AA, BB, CC;

        switch (this->get_cache()) {
        case LINAL_DEVICE_CACHE_NO_USE:
          AA = this->pull_obj_from_work( 1 );
          BB = this->pull_obj_from_work( 2 );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_ONLY:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_bin ( B );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_WRITE:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_bin ( B );
          CC = this->pull_obj_from_bin ( C );
          break;
        }

        Flat_Obj_ A_device, B_device, C_device;

        AA->extract( A_device, A.m, A.n );
        BB->extract( B_device, B.m, B.n );
        CC->extract( C_device, C.m, C.n );

        Int_ A_hit, B_hit, C_hit;

        A_device.set_matrix_from(A, A_hit);
        B_device.set_matrix_from(B, B_hit);
        C_device.set_matrix_from(C, C_hit);

        if (this->get_stat()) {
          this->_s_cache += ( A_hit + B_hit + C_hit );
          this->_s_device += 3;
          //std::cout << "gemm hit cache " << A_hit << B_hit << C_hit << std::endl;
        }

        AA->set_dirt(false);
        BB->set_dirt(false);
        CC->set_dirt(false);

        Int_ r_val = internal::gemm( transa, transb,
                                     alpha, A_device, B_device,
                                     beta, C_device );
        C_device.get_matrix_to(C);

        AA->set_use(false);
        BB->set_use(false);
        CC->set_use(false);

        return r_val;
      }
      virtual Int_ herk( Int_ uplo, Int_ trans,
                        FLA_Obj &alpha, FLA_Obj &A,
                        FLA_Obj &beta,  FLA_Obj &C ) {

        this->herk_stat();
	this->set_device();

        Flat_Obj AA, CC;

        switch (this->get_cache()) {
        case LINAL_DEVICE_CACHE_NO_USE:
          AA = this->pull_obj_from_work( 1 );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_ONLY:
          AA = this->pull_obj_from_bin ( A );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_WRITE:
          AA = this->pull_obj_from_bin ( A );
          CC = this->pull_obj_from_bin ( C );
          break;
        }

        Flat_Obj_ A_device, C_device;

        AA->extract( A_device, A.m, A.n );
        CC->extract( C_device, C.m, C.n );

        Int_ A_hit, C_hit;

        A_device.set_matrix_from(A, A_hit);
        C_device.set_matrix_from(C, C_hit);

        if (this->get_stat()) {
          this->_s_cache += ( A_hit + C_hit );
          this->_s_device += 2;
        }

        AA->set_dirt(false);
        CC->set_dirt(false);

        Int_ r_val = internal::herk( uplo, trans, 
                                    alpha, A_device, 
                                    beta, C_device );
        
        C_device.get_matrix_to(C);

        AA->set_use(false);
        CC->set_use(false);

        return r_val;
      }
      virtual Int_ trmm( Int_ side, Int_ uplo, Int_ trans, Int_ diag,
                        FLA_Obj &alpha, FLA_Obj &A, FLA_Obj &B ) {

        this->trmm_stat();
	this->set_device();
        
        Flat_Obj AA, BB;

        switch (this->get_cache()) {
        case LINAL_DEVICE_CACHE_NO_USE:
          AA = this->pull_obj_from_work( 1 );
          BB = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_ONLY:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_WRITE:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_bin ( B );
          break;
        }

        Flat_Obj_ A_device, B_device;

        AA->extract( A_device, A.m, A.n );
        BB->extract( B_device, B.m, B.n );

        Int_ A_hit, B_hit;

        A_device.set_matrix_from(A, A_hit);
        B_device.set_matrix_from(B, B_hit);

        if (this->get_stat()) {
          this->_s_cache += ( A_hit + B_hit );
          this->_s_device += 2;
        }

        AA->set_dirt(false);
        BB->set_dirt(false);

        Int_ r_val = internal::trmm( side, uplo, trans, diag, 
                                    alpha, A_device, B_device );

        B_device.get_matrix_to(B);

        AA->set_use(false);
        BB->set_use(false);

        return r_val;
      }
      virtual Int_ trsm( Int_ side, Int_ uplo, Int_ trans, Int_ diag,
                        FLA_Obj &alpha, FLA_Obj &A, FLA_Obj &B ) {

        this->trsm_stat(); 
	this->set_device();

        Flat_Obj AA, BB;

        switch (this->get_cache()) {
        case LINAL_DEVICE_CACHE_NO_USE:
          AA = this->pull_obj_from_work( 1 );
          BB = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_ONLY:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_WRITE:
          AA = this->pull_obj_from_bin ( A );
          BB = this->pull_obj_from_bin ( B );
          break;
        }

        Flat_Obj_ A_device, B_device;

        AA->extract( A_device, A.m, A.n );
        BB->extract( B_device, B.m, B.n );

        Int_ A_hit, B_hit;

        A_device.set_matrix_from(A, A_hit);
        B_device.set_matrix_from(B, B_hit);

        if (this->get_stat()) {
          this->_s_cache += ( A_hit + B_hit );
          this->_s_device += 2;
          //std::cout << "trsm hit cache " << A_hit << B_hit << std::endl;
        }

        AA->set_dirt(false);
        BB->set_dirt(false);

        Int_ r_val = internal::trsm( side, uplo, trans, diag, 
                                    alpha, A_device, B_device );

        B_device.get_matrix_to(B);

        AA->set_use(false);
        BB->set_use(false);

        return r_val;
      }

      virtual Int_ syrk( Int_ uplo, Int_ trans,
                        FLA_Obj &alpha, FLA_Obj &A,
                        FLA_Obj &beta,  FLA_Obj &C ) {

        this->syrk_stat();
	this->set_device();

        Flat_Obj AA, CC;

        switch (this->get_cache()) {
        case LINAL_DEVICE_CACHE_NO_USE:
          AA = this->pull_obj_from_work( 1 );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_ONLY:
          AA = this->pull_obj_from_bin ( A );
          CC = this->pull_obj_from_work( 0 );
          break;
        case LINAL_DEVICE_CACHE_READ_WRITE:
          AA = this->pull_obj_from_bin ( A );
          CC = this->pull_obj_from_bin ( C );
          break;
        }

        Flat_Obj_ A_device, C_device;

        AA->extract( A_device, A.m, A.n );
        CC->extract( C_device, C.m, C.n );

        Int_ A_hit, C_hit;

        A_device.set_matrix_from(A, A_hit);
        C_device.set_matrix_from(C, C_hit);

        if (this->get_stat()) {
          this->_s_cache += ( A_hit + C_hit );
          this->_s_device += 2;
        }

        AA->set_dirt(false);
        CC->set_dirt(false);

        Int_ r_val = internal::syrk( uplo, trans, 
                                    alpha, A_device, 
                                    beta, C_device );

        C_device.get_matrix_to(C);

        AA->set_use(false);
        CC->set_use(false);

        return r_val;
      }
    };
  }
}
#endif
